#!/usr/bin/env python
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
from std_msgs.msg import Float64MultiArray
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy

ARM_JOINT_NAMES = ["1", "2", "3", "4", "5"]
GRIPPER_JOINT_NAME = "6"

class LeaderFollowerBridge(Node):
    def __init__(self):
        super().__init__('so101_leader_follower_bridge')

        # ros2_control 控制器通常期望 RELIABLE QoS
        controller_qos = QoSProfile(
            depth=5,
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST,
            durability=DurabilityPolicy.VOLATILE
        )
        
        # 为 Leader 订阅者也指定一个匹配的 QoS (Reliable)
        leader_qos = QoSProfile(
            depth=10, # 匹配 publisher 的 10
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST,
            durability=DurabilityPolicy.VOLATILE
        )

        # 1. 订阅 Leader 臂发布的状态
        self.leader_state_sub = self.create_subscription(
            JointState,
            '/so101_leader/joint_states',
            self.leader_state_callback,
            leader_qos
        )

        # 2. 创建发布者, 指向 OPi 上的 ros2_control 控制器
        self.follower_arm_pub = self.create_publisher(
            Float64MultiArray,
            '/arm_controller/commands', # 发布给手臂控制器
            qos_profile=controller_qos
        )
        self.follower_gripper_pub = self.create_publisher(
            Float64MultiArray,
            '/gripper_controller/commands', # 发布给夹爪控制器
            qos_profile=controller_qos
        )

        self.get_logger().info("✅ Leader-Follower 桥接节点已启动。")
        self.get_logger().info("    订阅: /so101_leader/joint_states")
        self.get_logger().info("    发布: /arm_controller/commands, /gripper_controller/commands")
        # 调试日志
        self.get_logger().debug("--- 正在等待来自 /so101_leader/joint_states 的消息... ---")


    def leader_state_callback(self, msg: JointState):
        """接收 Leader 状态, 拆分, 并作为指令转发给 Follower"""
        
        # 确认回调被触发
        self.get_logger().debug(">>> [Bridge] 成功收到 Leader 状态消息!", throttle_duration_sec=2.0)
        
        try:
            positions = dict(zip(msg.name, msg.position))

            # 准备手臂指令 (Float64MultiArray)
            arm_msg = Float64MultiArray()
            # 按照 arm_controller 期望的顺序填充数据
            arm_msg.data = [positions[name] for name in ARM_JOINT_NAMES]

            # 准备夹爪指令 (Float64MultiArray)
            gripper_msg = Float64MultiArray()
            gripper_msg.data = [positions[GRIPPER_JOINT_NAME]]

            # 确认正在发布的数据
            self.get_logger().debug(f"    发布手臂: { [round(p, 2) for p in arm_msg.data] }", throttle_duration_sec=2.0)
            self.get_logger().debug(f"    发布夹爪: { [round(p, 2) for p in gripper_msg.data] }", throttle_duration_sec=2.0)

            # 发布指令
            self.follower_arm_pub.publish(arm_msg)
            self.follower_gripper_pub.publish(gripper_msg)

        except KeyError as e:
            # 如果 /so101_leader/joint_states 消息不完整, 打印警告
            self.get_logger().warn(f"在 /so101_leader/joint_states 消息中未找到关节 {e}", throttle_duration_sec=5.0)
        except Exception as e:
            self.get_logger().error(f"处理 Leader 状态时出错: {e}", throttle_duration_sec=5.0)


def main(args=None):
    rclpy.init(args=args)
    try:
        bridge_node = LeaderFollowerBridge()
        rclpy.spin(bridge_node)
    except KeyboardInterrupt:
        pass
    finally:
        if 'bridge_node' in locals() and rclpy.ok():
            bridge_node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()